# PPO expert policy can also drive in the safety-critical SafeMetaDrive environment

from metadrive import SafeMetaDriveEnv
from metadrive.examples import expert

# Set the environment config
config = {"start_seed": 1000,
          "num_scenarios": 1,
          "accident_prob": 1.0,
          "traffic_density": 0.25,
          }
env = SafeMetaDriveEnv(config=config)

print("Starting the environment ...")

ep_reward = 0.0
ep_cost = 0.0
obs, info = env.reset()

print(f'obs shape is {obs.shape}, obs type is {type(obs)}')
print(f"obs is {obs}")

frames = []

for i in range(1000):
    obs, reward, terminated, truncated, info = env.step(expert(env.agent))
    ep_reward += reward
    ep_cost += info["cost"]
    frame = env.render(mode="top_down", window=False, screen_size=(500, 500))
    frames.append(frame)
    if terminated or truncated:
        print(f"Arriving Destination: {info['arrive_dest']}")
        print(f"Episode reward: {ep_reward}.")
        print(f"Episode cost: {ep_cost}.")
        break

env.close()
print("MetaDrive successfully run!")
# render image
print("\nGenerate gif...")
import pygame
import numpy as np
from PIL import Image
import os

HOME = os.path.dirname(os.path.realpath(__file__))

imgs = [frame for frame in frames]
imgs = [Image.fromarray(img) for img in imgs]
imgs[0].save(f"{HOME}/demo.gif", save_all=True, append_images=imgs[1:], duration=50, loop=0)

